package com.evil.kingdom.handler;

import com.evil.kingdom.pojo.dto.alive.AliveMessage;
import com.evil.kingdom.pojo.entity.token.SocketInfo;
import com.evil.kingdom.utils.TokenUtil;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpHeaders;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.*;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author liyang
 * @date 2024-05-18 13:42
 */
@Slf4j
@Component
@AllArgsConstructor
public class MyWebSocketHandler implements WebSocketHandler {

    private static final ConcurrentHashMap<String, WebSocketSession> sessionMap = new ConcurrentHashMap<>();
    private final TokenUtil tokenUtil;

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws IOException {
        // 连接建立后的逻辑
        log.info("socket handler|After connection");
        if (null == session.getRemoteAddress()) {
            session.close();
        }
        // 记录socket ip
        String token = session.getHandshakeHeaders().getFirst(HttpHeaders.AUTHORIZATION);
        String hostAddress = session.getRemoteAddress().getAddress().getHostAddress();
        try {
            tokenUtil.setSocket(token, new SocketInfo(hostAddress));
            log.info("WebSocket set socket success| token:{}, hostAddress:{}", token, hostAddress);
        } catch (Exception e) {
            log.error("WebSocket set socket error| token:{}", token);
        }
        sessionMap.put(session.getId(), session);
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) {
        // 处理接收到的消息
        log.info("socket handler|Received message:{}", message.getPayload());
    }

    public void sendAll(AliveMessage<?> socketMessage) {
        sessionMap.forEach((k, v) -> {
            synchronized (v) {
                if (v.isOpen()) {
                    try {
                        v.sendMessage(new TextMessage(socketMessage.toMessageStr()));
                    } catch (IOException ignored) {
                    }
                }
            }
        });
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws IOException {
        // 处理传输错误
        log.info("socket handler|Transport error:{}", exception.getMessage());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws IOException {
        sessionMap.remove(session.getId());
        // 连接关闭后的逻辑
        log.info("socket handler|Connection closed:{}", closeStatus);
        if (null == session.getRemoteAddress()) {
            session.close();
        }
        // 记录socket ip
        String token = session.getHandshakeHeaders().getFirst(HttpHeaders.AUTHORIZATION);
        try {
            tokenUtil.detSocket(token);
            log.info("WebSocket delete socket success| token:{}", token);
        } catch (Exception e) {
            log.error("WebSocket delete socket error| token:{}", token);
        }
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }
}
